/*
Reference: 
    http://blog.roodo.com/rocksaying/archives/14848695.html
*/

#include <stdio.h>
#include <stdlib.h>  
#include <string.h>
#include <sys/types.h>
#include <sys/stat.h> 
#include <unistd.h> 
#include <sys/file.h> 
#include <sys/ipc.h> 
#include <sys/sem.h>
 
#if !defined(LOCK_SH) 
    #define LOCK_SH     1   /* shared lock */
#endif
#if !defined(LOCK_EX)
    #define LOCK_EX     2   /* exclusive lock */
#endif
#if !defined(LOCK_NB)
    #define LOCK_NB     4   /* don't block when locking */
#endif
#if !defined(LOCK_UN)
    #define LOCK_UN     8   /* unlock */
#endif
 
/* write by rock
  求 IPC 鍵值。
  須傳入一檔案名稱，若該檔不存在， IPC_key() 會自動建立。
 
  RC: 成功(鍵值)、失敗(-1)
  See also: stat(), ftok()
*/
static int IPC_key(const char*ipcname) 
{
    struct stat st;
    int fd;
    if(stat(ipcname, &st) == 0 && S_ISREG( st.st_mode)){
        return st.st_ino;
    }

    if((fd = creat(ipcname, 0644)) < 0) {
        return -1;
    }

    fstat(fd, &st);
    close(fd);
    return st.st_ino; // 以 i_node 的值做為 IPC key
                      // 不理會 SYSV 的 ftok() 函數。
                      // ftok() 是 SVR 所提供的函數，在 POSIX 及 BSD 中沒有。
                      // return ftok(ipcname,proj);
}
 
/*
  NAME: lock_shm() 鎖定共享記憶體。
 
    LOCK_SH 加上共享鎖定(讀取鎖定) 
    LOCK_EX 加上互斥鎖定(寫入鎖定)
    LOCK_NB 不擱置( non-blocking ，配合前兩者使用 )
    LOCK_UN 解除鎖定
   
  此函數提供了一個近似 flock() 行為的 shm 鎖定函數。
 
  RC: 成功(0)、失敗(-1) 
  See also: semget(), semctl(),  semop() 
*/
int lock_shm(const char* shmname, int op) 
{
    enum { EXLOCK,  SHLOCK} lock;
    struct sembuf lockop={0, 0, SEM_UNDO} /* sem 操作指令*/;
    int semkey,  semid;
    ushort sem_val[] = {
        1 /*Init value of EXLOCK*/, 
        0 /*Init value of SHLOCK*/ };
     
    //嘗試建立 sem 
    if((semkey = IPC_key(shmname)) < 0) { 
        return -1; 
    }
 
    if((semid = semget(semkey, 2, IPC_CREAT|IPC_EXCL|0640)) >= 0) { 
        // 初始 sem 的值 
        if(semctl(semid, 0, SETALL, sem_val) < 0) {  
            return -1; 
        }
    } 
    else { 
        //sem 可能建立了，取得semid  
        if((semid = semget(semkey, 1, 0)) < 0) { 
            return -1; 
        }
    } 
 
    /* 
        non-blocking 鎖定，只要多設一個旗標 IPC_NOWAIT 即達目的。 
        連傳回值也一樣：當無法立即取得對 sem 的操作時， semop() 不會被擱置，  
        會立即返回(傳回 -1 )，並設定 errno 為 EAGAIN 。
    */
    if( op & LOCK_NB ) {
        lockop.sem_flg |= IPC_NOWAIT;
    }
 
    if( op & LOCK_EX ) {   
        // 將 Exclusive Lock -1 ，讓其他 Process 暫停 
        lockop.sem_num = EXLOCK; 
        lockop.sem_op = -1; 
 
        if( semop(semid,  &lockop, 1) < 0 ) 
            return -1; 
                                                    
    
        // 將 Shared Lock 設為 0
        //printf("EXLOCK... wait SHLOCK be zero\n"); 
        lockop.sem_num  = SHLOCK; 
        lockop.sem_op = 0; 
        return semop(semid, &lockop, 1); 
    } 
    else if( op &  LOCK_SH ) { 
        // 取得 Shared Lock 前，要先拿到 Exclusive Lock ( -1 )
        // 避免資源互搶
        lockop.sem_num = EXLOCK; 
        lockop.sem_op = -1; 
        if( semop(semid, &lockop,  1) < 0 ) 
            return -1; 

        // 將 Shared Lock + 1
        //printf("EXLOCK... wait to get SHLOCK\n"); 
        lockop.sem_num  = SHLOCK; 
        lockop.sem_op = 1; 
        if( semop(semid, &lockop, 1) < 0 ) 
            return -1;

        // 將 Exclusive Lock + 1 釋放 Exclusive Lock 給其他 Process
        //printf("SHLOCK  ok\n"); 
        lockop.sem_num = EXLOCK; 
        lockop.sem_op = 1; 
        //printf("Release EXLOCK\n");  
        return semop(semid, &lockop, 1); 
    } 
 
    //以下動作皆為 LOCK_UN 
    if( semctl(semid, 0, GETALL,  sem_val) < 0 ) 
        return -1; 
    //printf("LOCK_UN, EXLOCK:%d, SHLOCK:%d\n",sem_val[0],  sem_val[1]); 
 
    /*
        如果沒有程序在等待取得鎖定的話，那解除鎖定的動作，將會是"刪除 sem "，而不是"改變 sem 的值"。 
        如果有程序等待取得 LOCK_EX 的話，那就是在等待 SHLOCK 的值變為 0 。
        如果有程序等待取得 LOCK_SH 的話，那就是在等待 EXLOCK 的值增加。 
    */ 
    if( semctl(semid, SHLOCK, GETZCNT, 0) == 0 && 
        semctl(semid, EXLOCK, GETNCNT, 0) == 0 ) 
    { 
        if( sem_val[EXLOCK] <= 1 && sem_val[SHLOCK]  == 0) { 
           // 沒有人再用 Exclusive/Shared Lock，刪除這個 Semphore
           /*
           printf("remove sem %d,Z:%d,N:%d\n", semid, 
           semctl(semid,SHLOCK,GETZCNT,0),
           semctl(semid,EXLOCK,GETNCNT,0));  
           */
           return semctl(semid, 0, IPC_RMID, 0); 
        } 
    } 
 
    if( sem_val[EXLOCK] > 0 && sem_val[SHLOCK]  > 0 ) {
        // 只有 Shared Lock 代表這個有 Shared Lock 的 Process Unlock 後，
        // 下一個等候 Shared Lock Unlock 的 Process 就可以繼續動作了。
        // 所以在此是 Unlock Shared Lock ( -1 )

        //printf("EXLOCK >0 and SHLOCK  >0: UNLOCK SHLOCK\n"); 
        lockop.sem_num = SHLOCK; 
        lockop.sem_op = -1; 
        return semop(semid, &lockop, 1); 
    } 
    else if( sem_val[EXLOCK]  <= 0 ) { 
        // 有 Exclusive/Shared Lock 代表有Process 上 Exclusive Lock / Shared Lock
        if( semctl(semid, SHLOCK, GETZCNT, NULL) >0 ) { 
            // 但有 Exclusive Lock 的 Process 需要等所有有 Shared Lock 的 Process
            // 都 Unlock 後才能 Unlock 並執行程序，所以在此先 Unlock Shared Lock ( -1 )

            //printf("UNLOCK SHLOCK\n"); 
            lockop.sem_num = SHLOCK;  
            lockop.sem_op = -1; 
        } 
        else { 
            // 所以有 Shared Lock 的 Process 都 Unlock 了
            // 代表正等待 Exclusive Lock 的 Process 可以起來動作了
            // 所以在此是將 Exclusive Lock 加為 1

            //printf("UNLOCK EXLOCK\n"); 
            lockop.sem_num = EXLOCK;  
            lockop.sem_op = 1; 
        } 
        return semop(semid, &lockop,  1); 
    } 
    return 0; 
}
